#!/usr/bin/env python
# coding:utf-8

import numpy as np
import cv2

from paddle.inference import Config
from paddle.inference import create_predictor


def resize(img, target_size):
    if not isinstance(img, np.ndarray):
        raise TypeError('image type is not numpy.')
    im_shape = img.shape
    im_scale_x = float(target_size) / float(im_shape[1])
    im_scale_y = float(target_size) / float(im_shape[0])
    img = cv2.resize(img, None, None, fx=im_scale_x, fy=im_scale_y)
    return img


def normalize(img, mean, std):
    img = img / 255.0
    mean = np.array(mean)[np.newaxis, np.newaxis, :]
    std = np.array(std)[np.newaxis, np.newaxis, :]
    img -= mean
    img /= std
    return img


# 图像预处理
def preprocess(img, img_size):
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    img = resize(img, img_size)
    img = img[:, :, ::-1].astype('float32')  # bgr -> rgb
    img = normalize(img, mean, std)
    img = img.transpose((2, 0, 1))  # hwc -> chw
    return img[np.newaxis, :]  #


# 后处理函数
def draw_bbox_image(frame, result, label_list, threshold=0.5):
    for res in result:
        cat_id, score, bbox = res[0], res[1], res[2:]
        if score < threshold or cat_id not in [10, 11]:
            continue
        for i in bbox:
            int(i)
        xmin, ymin, xmax, ymax = [int(pts) for pts in bbox]
        cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (255, 0, 255), 2)
        print('category id is {}, bbox is {}'.format(cat_id, bbox))
        try:
            label_id = label_list[int(cat_id)]
            # #cv2.putText(图像, 文字, (x, y), 字体, 大小, (b, g, r), 宽度)
            cv2.putText(frame, label_id, (int(xmin + 50), int(ymin - 2)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
            cv2.putText(frame, str(round(score, 2)), (int(xmin), int(ymin - 2)), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                        (0, 255, 0), 2)
        except KeyError:
            pass


def init_predictor(model_file,
                   params_file,
                   gpu_id =0):
    # 加载inference模型文件
    config = Config(model_file, params_file)
    config.enable_use_gpu(500, gpu_id)
    config.switch_ir_optim()
    config.enable_memory_optim()
    # Create predictor
    predictor = create_predictor(config)
    return predictor


def yolo_inference(predictor, img):
    # Set input
    im_size = 608
    data = preprocess(img, im_size)

    scale_factor = np.array([im_size * 1. / img.shape[0], im_size * 1. / img.shape[1]]).reshape((1, 2)).astype(
        np.float32)

    im_shape = np.array([im_size, im_size]).reshape((1, 2)).astype(np.float32)
    label_list = ['Youtong', 'Shu', 'Zhansunzhuangbei', 'Weiqiang', 'Henggan', 'Dianxiangan', 'Juma', 'Gongmen',
                  'Tongxinta', 'Langan', 'Che', 'Ren', 'Diaobao', 'Ludeng', 'Zujuehao', 'YewaiLukou', 'Danyaoxiang',
                  'Fangzi', 'ChengshiLukou', 'Tiesiwang', 'Sanjiaozhui', 'Zhangpeng']
    # label_list = ['Che', 'Ren']
    ims = [im_shape, data, scale_factor]
    input_names = predictor.get_input_names()
    for i, name in enumerate(input_names):
        input_tensor = predictor.get_input_handle(name)
        input_tensor.reshape(ims[i].shape)
        input_tensor.copy_from_cpu(ims[i].copy())
    # 执行Predictor
    predictor.run()
    # 获取输出
    results = []
    # 获取输出
    output_names = predictor.get_output_names()
    for i, name in enumerate(output_names):
        output_tensor = predictor.get_output_handle(name)
        output_data = output_tensor.copy_to_cpu()
        results.append(output_data)

    draw_bbox_image(img, results[0], label_list, threshold=0.3)
    return results[0]
